// ++++++++++++++++++++++++++++++++++++++++
// 《零基础Go语言算法实战》源码
// ++++++++++++++++++++++++++++++++++++++++
// Author:廖显东（ShirDon）
// Blog:https://www.shirdon.com/
// Gitee:https://gitee.com/shirdonl/goAlgorithms.git
// Buy link :https://item.jd.com/14101229.html
// ++++++++++++++++++++++++++++++++++++++++

package main

import (
	"fmt"
	"sort"
)

func findCriticalAndPseudoCriticalEdges(n int, edges [][]int) [][]int {
	// 将edges中每个切片的索引添加到该切片
	for i := range edges {
		edges[i] = append(edges[i], i)
	}
	sort.Slice(edges, func(i, j int) bool {
		return edges[i][2] < edges[j][2]
	})

	// 使用 Kruskal 查找 MST
	minWeight := mstWeight(n, edges, -1, -1)

	critical := make([]int, 0)
	pseudoCritical := make([]int, 0)

	for i := range edges {

		// 检查排除这条边的权重
		wgt := mstWeight(n, edges, i, -1)

		// 如果权重为-1（无法形成 MST）或权重增加边，则必须包含在 MST 中
		if wgt == -1 || wgt > minWeight {
			critical = append(critical, edges[i][3])
			continue
		}

		// 如果 weight == minWeight，则不确定这条边是否被排除
		wgt = mstWeight(n, edges, -1, i)
		if wgt == minWeight {
			pseudoCritical = append(pseudoCritical, edges[i][3])
		}
	}

	// 排序并返回
	sort.Ints(critical)
	sort.Ints(pseudoCritical)
	return [][]int{critical, pseudoCritical}
}

func mstWeight(n int, edges [][]int, skipIdx int, forceIncludeIndex int) int {
	dsu := NewUnionSet(n)

	var weight int
	if forceIncludeIndex != -1 {
		e := edges[forceIncludeIndex]
		dsu.union(e[0], e[1])
		weight += e[2]
	}

	for i, edge := range edges {
		if i == skipIdx {
			continue
		}
		a, b := edge[0], edge[1]
		ra, rb := dsu.find(a), dsu.find(b)
		// 当已经是同一组件的一部分时跳过
		if ra == rb {
			continue
		}
		dsu.union(a, b)
		weight += edge[2]
	}

	// 验证所有边都属于同一个组件
	a := dsu.find(0)
	for i := range dsu.parent[1:] {
		if dsu.find(1+i) != a {
			return -1
		}
	}
	return weight
}

type UnionSet struct {
	parent []int
	size   []int
}

func NewUnionSet(n int) *UnionSet {
	dsu := &UnionSet{
		parent: make([]int, n),
		size:   make([]int, n),
	}
	for i := 0; i < n; i++ {
		dsu.parent[i] = i
		dsu.size[i] = 1
	}
	return dsu
}

func (d *UnionSet) find(a int) int {
	if d.parent[a] == a {
		return a
	}
	root := d.find(d.parent[a])
	d.parent[a] = root
	return root
}

func (d *UnionSet) union(a, b int) {
	a = d.find(a)
	b = d.find(b)
	if a != b {
		if d.size[a] < d.size[b] {
			a, b = b, a
		}
		d.parent[b] = a
		d.size[a] += d.size[b]
	}
}

func main() {
	n := 5
	edges := [][]int{{0, 1, 1}, {1, 2, 1}, {2, 3, 2}, {0, 3, 2}, {0, 4, 3}, {3, 4, 3}, {1, 4, 6}}
	res := findCriticalAndPseudoCriticalEdges(n, edges)
	fmt.Println(res)
}

//$ go run interview8-2.go
//[[0 1] [2 3 4 5]]
